import matplotlib.pyplot as plt
import torch
from torch import nn

# =================== 曲线图 =========================
x = torch.linspace(-1, 1, 20)  # (20,)
# =================== 调整训练的参数格式 =========================
# 注意：放入线性模型中x的形态：（总的数量，输入特征的数量）
# 比如：当前的x总数量是20，而模型中输入的特征的数量为1 (20,1)
#      由于y是因变量，因此y的形状也被改为（20，1）
x = x.reshape(20, 1)

y = x ** 2
plt.plot(x, y, 'ro')

# =================== 构建模型 =========================
# 注意：构建模型需要 torch.nn
# Sequential 组装模型的对象
model = nn.Sequential(
    # 输入
    nn.Linear(1, 10),  # 10是隐藏层的神经元数量
    # 激活函数
    nn.Tanh(),
    # 输出
    nn.Linear(10, 1)
)
# =================== 定义模型训练的配置 =========================
# 损失函数
criterion = nn.MSELoss()  # MSELoss 均方差损失（总距离公式）
# 优化器
# model.parameters() 实际上模型所有的w,b
# lr => learning-rate => 学习率
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
# =================== 训练模型 =========================
# 模型训练次数
epochs = 10000
# 开始训练
for epoch in range(epochs):
    # 优化器将所有的w，b的导数(专业词汇：梯度)清零
    optimizer.zero_grad()
    # 预测
    y_predict = model(x)
    # 计算损失
    loss = criterion(y_predict, y)  # 前向传播
    loss.backward()  # 反向传播
    # 更新w、b的参数
    optimizer.step()

    print(f"{epoch + 1} / {epochs} -- loss:{loss.item():.4f}")  # :.4f保留四位小数
# =================== 绘制预测线 =========================
# 开启模型的预测模式（评估模式）
model.eval()
y_predict = model(x)
plt.plot(x.detach().numpy(), y_predict.detach().numpy(), 'b--')
plt.show()
